package task

import (
	"context"
	"fmt"
	"sync"
	"time"
	"xagent/internal/memory"
	"xagent/internal/scheduler"
	proto "xagent/proto/v1"
)

// TaskManager manages task lifecycle and scheduling
type TaskManager struct {
	tasks          map[string]*Task
	processor      *TaskProcessor
	memoryManager  memory.MemoryManager
	statusHandlers map[proto.TaskStatus][]func(task *Task)
	mutex          sync.RWMutex
	ctx            context.Context
	cancel         context.CancelFunc
	queue          []*Task
	maxRetries     int
	retryDelay     time.Duration
	scheduler      *scheduler.Scheduler
}

// NewTaskManager creates a new task manager instance
func NewTaskManager(memoryManager memory.MemoryManager) *TaskManager {
	ctx, cancel := context.WithCancel(context.Background())
	tm := &TaskManager{
		tasks:          make(map[string]*Task),
		memoryManager:  memoryManager,
		statusHandlers: make(map[proto.TaskStatus][]func(task *Task)),
		ctx:            ctx,
		cancel:         cancel,
		queue:          make([]*Task, 0),
		maxRetries:     3,
		retryDelay:     time.Second * 5,
		scheduler:      scheduler.NewScheduler(10, scheduler.NewResourcePool(100, 5, 0.8)),
	}
	tm.processor = NewTaskProcessor(tm, memoryManager, 1000)
	return tm
}

// AddTask adds a new task to the manager
func (tm *TaskManager) AddTask(task *Task) error {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	if task.ID == "" {
		return fmt.Errorf("task ID cannot be empty")
	}

	if _, exists := tm.tasks[task.ID]; exists {
		return fmt.Errorf("task with ID %s already exists", task.ID)
	}

	if task.Metadata == nil {
		task.Metadata = &proto.TaskMetadata{
			Metadata:          make(map[string]string),
			RequiredResources: make(map[string]float64),
		}
	}

	tm.tasks[task.ID] = task
	return nil
}

// GetTask retrieves a task by ID
func (tm *TaskManager) GetTask(taskID string) (*Task, error) {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()

	task, exists := tm.tasks[taskID]
	if !exists {
		return nil, fmt.Errorf("task with ID %s not found", taskID)
	}

	return task, nil
}

// UpdateTaskStatus updates the status of a task
func (tm *TaskManager) UpdateTaskStatus(taskID string, status proto.TaskStatus) error {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	task, exists := tm.tasks[taskID]
	if !exists {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	task.Status = status

	// Call status handlers
	if handlers, ok := tm.statusHandlers[status]; ok {
		for _, handler := range handlers {
			handler(task)
		}
	}

	return nil
}

// RegisterStatusHandler registers a handler for task status changes
func (tm *TaskManager) RegisterStatusHandler(status proto.TaskStatus, handler func(task *Task)) {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	if _, ok := tm.statusHandlers[status]; !ok {
		tm.statusHandlers[status] = make([]func(task *Task), 0)
	}

	tm.statusHandlers[status] = append(tm.statusHandlers[status], handler)
}

// GetTasksByStatus returns all tasks with the specified status
func (tm *TaskManager) GetTasksByStatus(status proto.TaskStatus) []*Task {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()

	var tasks []*Task
	for _, task := range tm.tasks {
		if task.Status == status {
			tasks = append(tasks, task)
		}
	}

	return tasks
}

// RemoveTask removes a task from the manager
func (tm *TaskManager) RemoveTask(taskID string) error {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	if _, exists := tm.tasks[taskID]; !exists {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	delete(tm.tasks, taskID)
	return nil
}

// Close cleans up resources
func (tm *TaskManager) Close() {
	tm.cancel()
}

// processTask handles task processing lifecycle
func (tm *TaskManager) processTask(task *Task) {
	// Process task
	err := tm.processor.ProcessTask(tm.ctx, task)
	if err != nil {
		tm.handleTaskError(task, err)
		return
	}

	// Update task completion
	tm.UpdateTaskStatus(task.ID, proto.TaskStatus_TASK_STATUS_COMPLETED)
}

// handleTaskError handles task processing errors
func (tm *TaskManager) handleTaskError(task *Task, err error) {
	// Update task status with error
	result := NewTaskResult()
	result.SetError(err)
	tm.UpdateTaskStatus(task.ID, proto.TaskStatus_TASK_STATUS_FAILED)
}

// GetAllTasks returns all managed tasks
func (tm *TaskManager) GetAllTasks() []*Task {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()

	tasks := make([]*Task, 0, len(tm.tasks))
	for _, task := range tm.tasks {
		tasks = append(tasks, task)
	}
	return tasks
}

// ClearCompletedTasks removes completed tasks from memory
func (tm *TaskManager) ClearCompletedTasks() {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	for id, task := range tm.tasks {
		if task.Status == proto.TaskStatus_TASK_STATUS_COMPLETED {
			delete(tm.tasks, id)
		}
	}
}

// sortQueue sorts the task queue by priority
func (tm *TaskManager) sortQueue() {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	// Sort by priority (higher priority first)
	for i := 0; i < len(tm.queue)-1; i++ {
		for j := i + 1; j < len(tm.queue); j++ {
			if tm.queue[j].Priority > tm.queue[i].Priority {
				tm.queue[i], tm.queue[j] = tm.queue[j], tm.queue[i]
			}
		}
	}
}
